
#include "nmkb.h"
NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath);
NTSTATUS ForwardAndForget(PDEVICE_OBJECT DeviceObject,PIRP Irp);
VOID OnUnload(PDRIVER_OBJECT DriverObject);
void uninstall(); void find_video();void find_kbd_call(PDEVICE_OBJECT);
NTSTATUS inspector(PDEVICE_OBJECT DeviceObject,PIRP Irp);
NTSTATUS inspector_sw(PDEVICE_OBJECT DeviceObject,PIRP Irp);
NTSTATUS m_addDevice(PDRIVER_OBJECT  DriverObject,PDEVICE_OBJECT  PhysicalDeviceObject );
PDEVICE_OBJECT find_dev(PDEVICE_OBJECT d,wchar_t* name);
void hook_it();
#pragma alloc_text(PAGE,ForwardAndForget)
#pragma alloc_text(PAGE,OnUnload)
#pragma alloc_text(PAGE,uninstall)
#pragma alloc_text(PAGE,m_addDevice)
#pragma alloc_text(PAGE,find_dev)
#pragma alloc_text(INIT,find_video)
#pragma alloc_text(INIT,find_kbd_call) 
#pragma alloc_text (INIT, hook_it)
#pragma alloc_text (INIT, DriverEntry)


static PDEVICE_OBJECT dev=0;
static PFILE_OBJECT kbdf=0;
static int installed = 0;
static PDEVICE_OBJECT video0 = 0;
static unsigned char code[] = {0x55, 0x8b, 0xec, 0x8b,0x45 ,0x1c ,0xc6 ,0x00 ,0x01 ,0xb0 ,0x01 ,0x5d ,0xc2 ,0x1c ,0x00 ,0xcc};
static wchar_t keybd[] = L"\\Device\\KeyboardClass0";
static wchar_t driver[] = L"\\Driver\\nv";/*ati2mtag*/
static wchar_t monitor[] = L"\\device\\videopdo2";
static wchar_t pointer[] = L"\\Device\\PointerClass0";
static unsigned offset = 0x58;
static UNICODE_STRING ukeybd,udriver,umonitor,upointer;
static int hook_mouse = 0;
RTL_QUERY_REGISTRY_TABLE qtable[]={
{0,RTL_QUERY_REGISTRY_DIRECT,L"keyboard",&ukeybd,REG_SZ,&keybd,0},
{0,RTL_QUERY_REGISTRY_DIRECT,L"driver",&udriver,REG_SZ,&driver,0},
{0,RTL_QUERY_REGISTRY_DIRECT,L"monitor",&umonitor,REG_SZ,&monitor,0},
{0,RTL_QUERY_REGISTRY_DIRECT,L"mouse",&upointer,REG_SZ,&pointer,0},
{0,RTL_QUERY_REGISTRY_DIRECT,L"hookmouse",&hook_mouse,REG_DWORD,0,0},
{0,RTL_QUERY_REGISTRY_DIRECT,L"hidoffset",&offset,REG_DWORD,(void*)0x58,0},
{0}
};
KEYBOARD_INPUT_DATA copycode[] = {
    {0,0x1D,KEY_MAKE,0,0},
    {0,0x2E,KEY_MAKE,0,0},
    {0,0x2E,KEY_BREAK,0,0},
    {0,0x1d,KEY_BREAK,0,0},
    {0}
};
KEYBOARD_INPUT_DATA pastecode[] = {
    {0,0x1D,KEY_MAKE,0,0},
    {0,0x2F,KEY_MAKE,0,0},
    {0,0x2F,KEY_BREAK,0,0},
    {0,0x1d,KEY_BREAK,0,0},
    {0}
};

KEYBOARD_INPUT_DATA extra_code[] = {
    {0,0x76,KEY_MAKE,0,0},
    {0,0x6E,KEY_MAKE,0,0},
    {0,0x76,KEY_BREAK,0,0},
    {0,0x6E,KEY_BREAK,0,0},
    {0}
};
static WCHAR   deviceLinkBuffer[]  = L"\\DosDevices\\nmkb";
static WCHAR   deviceNameBuffer[]  = L"\\Device\\nmkb";

static kbd_call_t kbd_call;
static PDEVICE_OBJECT kbd_dev;
static mou_call_t mou_call;
static PDEVICE_OBJECT blu_dev = 0;
static PDEVICE_OBJECT* orig_blu = 0;
static PDEVICE_OBJECT mydev=0;
static PDEVICE_OBJECT odev = 0;
static long* orig_serv = 0;
static PKEVENT unload_ev = 0;
static AddDevice_t mouclass_add = 0;
static AddDevice_t* orig_add = 0;
VOID NTAPI set_sus ( int c);
void uninstall(){
    void * cc = ExAllocatePool(NonPagedPool,sizeof(code));
    if(cc) {
       PIRP reqn;
       INTERNAL_I8042_HOOK_KEYBOARD buf={0,0,cc};
       RtlCopyMemory(cc,code,sizeof(code));
       reqn = IoBuildDeviceIoControlRequest(IOCTL_INTERNAL_I8042_HOOK_KEYBOARD,dev,&buf,sizeof(buf),0,0,TRUE,0,0);
       if(dev && installed && reqn) 
            IoCallDriver(dev,reqn);
    }

}
VOID OnUnload(PDRIVER_OBJECT DriverObject)
{
   uninstall();
   DPrint(("nmkb: OnUnload 1\n"));
   if(kbdf) {
       DPrint(("nmkb: deref %x\n",kbdf));
       ObDereferenceObject(kbdf);    
   }
   if (orig_serv) {
       *orig_serv = (long)mou_call;
   }
   if(hook_mouse && orig_blu){
       KEVENT ev;
       LARGE_INTEGER ll;
       ll.QuadPart = 50000000;
       KeInitializeEvent(&ev,NotificationEvent,0);
       unload_ev = &ev;
       KeWaitForSingleObject(&ev,Executive,KernelMode,0,0);
       *orig_blu = blu_dev;
   }
   if(orig_add) {
        *orig_add = mouclass_add;
   }
   if(mydev){
       IoDeleteDevice(mydev);
   }
   if(odev) {
       UNICODE_STRING deviceLinkUnicodeString;
       RtlInitUnicodeString( &deviceLinkUnicodeString, deviceLinkBuffer );
       IoDeleteSymbolicLink(&deviceLinkUnicodeString);
       IoDeleteDevice(odev);
   }
}

NTSTATUS ForwardAndForget(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
   Irp->IoStatus.Status = STATUS_SUCCESS;
   IoCompleteRequest(Irp, IO_NO_INCREMENT);
   return Irp->IoStatus.Status;
}
BOOLEAN
KbFilter_IsrHook(
    PDEVICE_OBJECT         DeviceObject,               
    PKEYBOARD_INPUT_DATA   CurrentInput, 
    POUTPUT_PACKET         CurrentOutput,
    UCHAR                  StatusByte,
    PUCHAR                 DataByte,
    PBOOLEAN               ContinueProcessing,
    PKEYBOARD_SCAN_STATE   ScanState
    )
{

    UCHAR scan = *DataByte;
    USHORT status = CurrentInput->Flags;
	static isdat = 0;
	static  unsigned _int64 tick = 0;
	static ison=0;
    //DPrint(("filter_key: %x %x\n",status,scan));

    //BOOLEAN ret = KeInsertQueueDpc(&dcall,(PVOID)scan,(PVOID)status);
    //if(!ret) {
    //    DbgPrint("!!!!!!!not queued");
    //}
#define ctrl 0x9d 
#define right_flow    0xe9
    if(scan == ctrl && (status & KEY_E0)) {
		unsigned _int64 curr = KeQueryInterruptTime();
		if(curr-tick < 5000000i64){
			isdat++;
		}
		else {
			isdat = 0;
		}

		tick = curr;
	    if(isdat >= 2) {
		    isdat = 0;
            set_sus(++ison%2);
        }
    }
    *ContinueProcessing = TRUE;
    return TRUE;
}

VOID NTAPI find_video()
{
   NTSTATUS Status;
   PDRIVER_OBJECT ati;
   PDEVICE_OBJECT p;
   wchar_t ss[128];
   OBJECT_NAME_INFORMATION ni= {{0,sizeof(ss),ss}};
   
   Status = ObReferenceObjectByName(&udriver,OBJ_CASE_INSENSITIVE,NULL,0,IoDriverObjectType,KernelMode,NULL,(PVOID*)&ati);
   if(Status != STATUS_SUCCESS) {   DbgPrint("Failed video0 %x\n",Status); return;}

   p = ati->DeviceObject;
	while (p)
	{
		ULONG l = 0;
		Status = ObQueryNameString(p,&ni,128,&l);
		if (Status == STATUS_SUCCESS && !_wcsicmp(umonitor.Buffer,ni.Name.Buffer)) /*videopdo1*/
		{
			DPrint(("Find video0!!!\n"));
			video0 = p;
            break;
		}
		p = p->NextDevice;
	}
   ObDereferenceObject(ati);
}
VOID NTAPI set_sus ( int c)
{
    POWER_STATE ps;
    ULONG state;
    //DbgPrint("In set_sus IRQL: %d\n",KeGetCurrentIrql());
	if(!video0) {
         DbgPrint("video0 not ready");
         return;
	}
    state = video0->DeviceObjectExtension->PowerFlags;
    DbgPrint("%x",state);
	if( (state & 0xff) == 0x10)
	{
		ps.DeviceState = PowerDeviceD3;
		ps.SystemState=PowerSystemSleeping3;
	}
	else {
		ps.DeviceState = PowerDeviceD0;
		ps.SystemState=PowerSystemWorking;
	}
    PoRequestPowerIrp(video0,IRP_MN_SET_POWER,ps,0,0,0);

}
VOID mycall(
    IN PKDPC Dpc,
    IN PVOID DeferredContext,
    IN PVOID a,
    IN PVOID b
    )
{
    DbgPrint("mycall: %x %x",a,b);
}

int  query_name(PVOID p, PWSTR s) {
    wchar_t ss[128]; ULONG len;
    OBJECT_NAME_INFORMATION ni= {{0,sizeof(ss),ss}};
    return (ObQueryNameString(p,&ni,128,&len) == STATUS_SUCCESS && !_wcsicmp(ni.Name.Buffer,s));
}
PDEVICE_OBJECT find_dev(PDEVICE_OBJECT d,wchar_t* name) {
    PDEVICE_OBJECT p = d;
    wchar_t ss[128]; ULONG len;
    OBJECT_NAME_INFORMATION ni= {{0,sizeof(ss),ss}};
    while (p)
    {
        PDRIVER_OBJECT dr = p->DriverObject;
        if (dr) {
            if ( ObQueryNameString(dr,&ni,128,&len) == STATUS_SUCCESS && !_wcsicmp(ni.Name.Buffer,name)) { 
                return p;
            }
        }
        p = ((PMDEVOBJ_EXTENSION)(p->DeviceObjectExtension))->AttachedTo;
    }
    p = d;
    while (p)
    {
        PDRIVER_OBJECT dr = p->DriverObject;
        if (dr) {
            if ( ObQueryNameString(dr,&ni,128,&len) == STATUS_SUCCESS && !_wcsicmp(ni.Name.Buffer,name)) { 
                return p;
            }
        }
        p = p->AttachedDevice;
    }

    return 0;
}

void find_kbd_call(PDEVICE_OBJECT d) {
    
    PDEVICE_OBJECT p = find_dev(d,L"\\driver\\i8042prt");
    if (p) {
        PCOMMON_DATA ext = p->DeviceExtension ;
        kbd_call = (kbd_call_t)ext->ConnectData.ClassService;
        kbd_dev = ext->ConnectData.ClassDeviceObject;
        DPrint(("service found: %x %x %x",kbd_call,ext->ConnectData.ClassDeviceObject,d));
    }

}
VOID __stdcall  mouse_callback(
    IN PDEVICE_OBJECT  DeviceObject,
    IN PMOUSE_INPUT_DATA  InputDataStart,
    IN PMOUSE_INPUT_DATA  InputDataEnd,
    IN OUT PULONG  InputDataConsumed
    )
{
    PMOUSE_INPUT_DATA d = InputDataStart;
    if (d->ButtonFlags) {
        DPrint(("flag: %x",d->ButtonFlags));
    }
    mou_call(DeviceObject,InputDataStart,InputDataEnd,InputDataConsumed);
}
void hook_mouhid(PDEVICE_OBJECT t)
{
        PVOID ext = t->DeviceExtension;
    /*mou_call = (mou_call_t)(*(PLONG)((PUCHAR)ext+offset));
    DPrint(("mouse service found: %x",mou_call));
    orig_serv = (long*)((PUCHAR)ext+offset);
    *orig_serv = (long)&mouse_callback;*/
    orig_blu = (PDEVICE_OBJECT*)((PUCHAR)ext+4);
    blu_dev = *orig_blu;
    *orig_blu = mydev;
    DPrint(("Hook success dev:%x",blu_dev));
    DPrint(("Hook success new:%x",*orig_blu));                    

}
void hook_it() {
    NTSTATUS r = 0;
    PDRIVER_OBJECT mou=0;
    UNICODE_STRING mouclass;
    RtlInitUnicodeString(&mouclass,L"\\Driver\\Mouclass");
    r = ObReferenceObjectByName(&mouclass,OBJ_CASE_INSENSITIVE,0,0,IoDriverObjectType,KernelMode,0,&mou);
    DPrint(("!!!!!!!!!%x",r));
    orig_add = &mou->DriverExtension->AddDevice;
    mouclass_add = mou->DriverExtension->AddDevice;
    mou->DriverExtension->AddDevice = m_addDevice;
    DPrint(("AddDevice orig,%x",mouclass_add));
    DPrint(("AddDevice new, %x",*orig_add));

    if (r == STATUS_SUCCESS) {
        PDEVICE_OBJECT p = mou->DeviceObject;
        while (p) {
                PDEVICE_OBJECT t = 0;
                if (find_dev(p,L"\\Driver\\HidBth") && (t = find_dev(p,L"\\Driver\\Mouhid")) ) {
                    hook_mouhid(t);
                }
                break;
            p = p->NextDevice;
        }

    }
    
    if(mou) {
        ObDereferenceObject(mou);
    }
}
NTSTATUS m_addDevice(PDRIVER_OBJECT  DriverObject,PDEVICE_OBJECT  pdo )
{
    NTSTATUS ret = 0;
    PDEVICE_OBJECT p = 0;
    DPrint(("Add device %x",pdo));
    if (find_dev(pdo,L"\\Driver\\HidBth") && (p = find_dev(pdo,L"\\Driver\\Mouhid"))) {
        hook_mouhid(p);
    }
    ret = mouclass_add(DriverObject,pdo);
    return ret;
}


NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegistryPath)
{
    NTSTATUS ret = 0;
    PIRP req = 0;int i = 0;
    UNICODE_STRING          deviceNameUnicodeString;
    UNICODE_STRING          deviceLinkUnicodeString;    

    INTERNAL_I8042_HOOK_KEYBOARD buffer={0,0,KbFilter_IsrHook};
    ret = RtlQueryRegistryValues(RTL_REGISTRY_ABSOLUTE,RegistryPath->Buffer,qtable,0,0);
    if(ret != STATUS_SUCCESS) {
        DPrint(("query value failed: %S %x\n",RegistryPath->Buffer,ret));
        return STATUS_EA_CORRUPT_ERROR;
    }

    ret = IoGetDeviceObjectPointer(&ukeybd,MAXIMUM_ALLOWED,&kbdf,&dev);
    if(ret != STATUS_SUCCESS) {
        DPrint(("keyboard not found: %S\n",ukeybd.Buffer));
        return ret;
    }
    DPrint(("nmkb: DriverEntry %d %x\n",ret,dev));

    find_video();
    if(video0 == 0 ){
        DPrint(("video not found: %S %S\n",udriver.Buffer,umonitor.Buffer));
        return STATUS_DEVICE_NOT_CONNECTED; 
    }


    if(hook_mouse) {
        find_kbd_call(dev); 
    }
    
    //KeInitializeDpc(&dcall,mycall,0);

    req = IoBuildDeviceIoControlRequest(IOCTL_INTERNAL_I8042_HOOK_KEYBOARD,dev,&buffer,sizeof(buffer),0,0,TRUE,0,0);

    if (req == 0){
        DPrint(("nmkb: hook failed\n"));
        return STATUS_DEVICE_NOT_CONNECTED;
    }
    ret = IoCallDriver(dev,req);
    if(ret != STATUS_SUCCESS) 
        return ret;
    installed = 1;

    RtlInitUnicodeString (&deviceNameUnicodeString,
                          deviceNameBuffer );
    RtlInitUnicodeString (&deviceLinkUnicodeString,
                          deviceLinkBuffer );


    ret = IoCreateDevice(DriverObject,0,&deviceNameUnicodeString,FILE_DEVICE_UNKNOWN,0,0,&odev);
    if( !NT_SUCCESS(ret)) {uninstall();return ret;}
    //odev->Flags |= DO_BUFFERED_IO;
    ret = IoCreateSymbolicLink (&deviceLinkUnicodeString,&deviceNameUnicodeString );
    if( !NT_SUCCESS(ret)) {uninstall();IoDeleteDevice(odev);return ret;}

    if (hook_mouse) {
        ret = IoCreateDevice(DriverObject,0,0,FILE_DEVICE_UNKNOWN,0,0,&mydev);
        if(NT_SUCCESS(ret)){
            hook_it();
        }
    }

    for (i = 0; i <= IRP_MJ_MAXIMUM_FUNCTION; i++) {
        DriverObject->MajorFunction[i] = inspector_sw;
    }
    DriverObject->DriverUnload = OnUnload;
    return STATUS_SUCCESS;
}
NTSTATUS blu_complete(PDEVICE_OBJECT  DeviceObject,PIRP  Irp,PVOID  Context)
{
    PIO_STACK_LOCATION irpStack  = IoGetCurrentIrpStackLocation(Irp);
    ULONG len = irpStack->Parameters.Read.Length;
    PMDL mdl = Irp->MdlAddress;
    static ULONG left = 0;
    static ULONG right = 0;
    

    if (mdl && len == 7){
        PUCHAR base = MmGetMdlVirtualAddress(mdl);
        UCHAR code;
        if (!base) return STATUS_SUCCESS;
        code = base[6];
        DPrint(("oncall: %02X %02X %02X %02X %02X %02X %02X",base[0],base[1],base[2],base[3],base[4],base[5],base[6]));
        //DPrint(("%02X",code));
        if (code == 0xff) {
            left +=1;
            //DPrint(("%x",code));
            if (left == 1) {
                kbd_call(kbd_dev,&copycode[0],&copycode[1], &len);
            }
            else if (left == 2) {
                kbd_call(kbd_dev,&copycode[3],&copycode[4], &len);
            }
            if (left > 1) {
                 kbd_call(kbd_dev,&extra_code[0],&extra_code[1], &len);
            }
        }
        else if (code == 1) {
            right += 1;
            if (right == 1) {
                kbd_call(kbd_dev,&pastecode[0],&pastecode[1], &len);
            }
            if (right == 2) {
                kbd_call(kbd_dev,&pastecode[3],&pastecode[4], &len);
            }
            if (right > 1) {
                kbd_call(kbd_dev,&extra_code[1],&extra_code[2], &len);
            }
        }
        else {
            if (left == 1){
                kbd_call(kbd_dev,&copycode[1],&copycode[4], &len);
            }
            else if (left > 1) {
                kbd_call(kbd_dev,&extra_code[2],&extra_code[3], &len);
            }
            else if (right == 1) {
                kbd_call(kbd_dev,&pastecode[1],&pastecode[4], &len);
            }
            else if (right > 1) {
                kbd_call(kbd_dev,&extra_code[3],&extra_code[4], &len);
            }
            left = right = 0;
        }
    }
    return STATUS_SUCCESS;
}
#define MY_CODE(x) CTL_CODE(FILE_DEVICE_UNKNOWN,x,METHOD_BUFFERED,FILE_ANY_ACCESS)

NTSTATUS inspector_sw_odev(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
    PIO_STACK_LOCATION irpStack = IoGetCurrentIrpStackLocation(Irp);
    DPrint(("Some Driver Call %x",irpStack->MajorFunction));
    switch(irpStack->MajorFunction) {
    case IRP_MJ_DEVICE_CONTROL:
        DPrint(("IRP_MJ_DEVICE_CONTROL\n"));
        {

        PVOID inputBuffer             = Irp->AssociatedIrp.SystemBuffer;
        ULONG inputBufferLength       = irpStack->Parameters.DeviceIoControl.InputBufferLength;
        PVOID outputBuffer            = Irp->AssociatedIrp.SystemBuffer;
        ULONG outputBufferLength      = irpStack->Parameters.DeviceIoControl.OutputBufferLength;
        ULONG ioControlCode           = irpStack->Parameters.DeviceIoControl.IoControlCode;

        DPrint(("Control Code: %x",ioControlCode));
        if (ioControlCode == 0x333) {
            PoSetSystemState(ES_USER_PRESENT);
        }
        else if (ioControlCode == 0x334) {
            set_sus(1);
        }
        else if (ioControlCode == 0x335) {
            PIRP q = IoAllocateIrp(8,0);
            PIO_STACK_LOCATION n = IoGetNextIrpStackLocation(q);
            NTSTATUS r = 0;
            POWER_SEQUENCE sec = {0};
            if  (q == 0) {
                break;
            }
            RtlZeroMemory(n,sizeof(IO_STACK_LOCATION));
            n->CompletionRoutine = 0;

            n->MajorFunction = IRP_MJ_POWER;
            n->MinorFunction = IRP_MN_POWER_SEQUENCE;
            n->Parameters.PowerSequence.PowerSequence = &sec;
            q->RequestorMode = KernelMode;
            q->Flags =  
            r = PoCallDriver(video0,q);
            if (STATUS_SUCCESS == r && q->IoStatus.Status == STATUS_SUCCESS) {
               // DPrint(("%d %d %d",n->Parameters.PowerSequence.PowerSequence->SequenceD1,n->Parameters.PowerSequence.PowerSequence->SequenceD2,n->Parameters.PowerSequence.PowerSequence->SequenceD3));
            }
            else {
                DPrint(("Failed:%x",r));
            }

        }
        else if (ioControlCode == MY_CODE(0x336)) {
            void net_con(PVOID input,ULONG isize,PVOID output, ULONG osize,PIRP Irp);
            net_con(inputBuffer,inputBufferLength,outputBuffer,outputBufferLength,Irp);
        }
        }

     case IRP_MJ_CREATE:
     default:
            ;
    }
    Irp->IoStatus.Status = STATUS_SUCCESS;
    IoCompleteRequest( Irp, IO_NO_INCREMENT );
    return STATUS_SUCCESS;
}
NTSTATUS inspector_sw(PDEVICE_OBJECT DeviceObject,PIRP Irp)
{
    PIO_STACK_LOCATION irpStack = IoGetCurrentIrpStackLocation(Irp);
    if (DeviceObject == odev) return inspector_sw_odev(DeviceObject,Irp);
    switch(irpStack->MajorFunction) {
    case IRP_MJ_READ:
        if (unload_ev) {
            KeSetEvent(unload_ev,0,0);
            break;
        } else {
            IoCopyCurrentIrpStackLocationToNext(Irp);
            IoSetCompletionRoutine(Irp,blu_complete,0,1,0,0);
            return IoCallDriver(blu_dev,Irp);
        }
    case IRP_MJ_POWER:
        DPrint(("Power off"));  
        PoStartNextPowerIrp(Irp);
        IoSkipCurrentIrpStackLocation(Irp);
        return PoCallDriver(blu_dev,Irp);
    case IRP_MJ_PNP:
        if (irpStack->MinorFunction == IRP_MN_REMOVE_DEVICE) {
            DPrint(("Removal!!"));
            *orig_blu = blu_dev;
            orig_blu = 0;
        }
        break;
    }
    IoSkipCurrentIrpStackLocation(Irp);
    return IoCallDriver(blu_dev, Irp);
} 
